import os
import cv2
import fire
import glob
import time
import logging
import numpy as np
import xml.dom.minidom as minidom
from scipy.signal import argrelmin, argrelmax


def getimg(img_path):
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)
    img_name = os.path.basename(img_path).split(r'.')[0]
    return img, img_name


def cut_by_grid(img_data, rows, cols):
    # 对比较正的层前图，利用栅线信息进行边缘切割
    hh, ww = img_data.shape[:2]
    col_lines, row_lines = get_cut_line_none_precut_edge(img_data, rows, cols)
    if abs(col_lines[cols // 2] - ww //
           2) / ww > 0.01 or abs(row_lines[rows // 2] - hh // 2) / hh > 0.01:
        col_lines, row_lines = get_cut_line_none_precut_edge(
            img_data, rows, cols)
    # print(col_lines, row_lines)
    col_step = (col_lines[-2] - col_lines[1]) // (cols - 2)
    row_step = (row_lines[-2] - row_lines[1]) // (rows - 2)
    xmin, xmax = col_lines[1] - col_step, col_lines[-2] + col_step
    ymin, ymax = row_lines[1] - row_step, row_lines[-2] + row_step
    tl, tr, bl, br = (xmin, ymin), (xmax, ymin), (xmin, ymax), (xmax, ymax)
    return tl, tr, bl, br


def get_precut(img_data, rows, cols):
    '''
    '''
    offset = 10
    hh, ww = img_data.shape[:2]

    if rows > 4 and cols > 4:
        # 对比较正的层前图，利用栅线信息进行边缘切割
        #             try:
        tl, tr, bl, br = cut_by_grid(img_data, rows, cols)
        min_x = max(min(tl[0], bl[0]) - offset, 0)
        max_x = min(max(tr[0], br[0]) + offset, ww - 1)
        min_y = max(min(tl[1], tr[1]) - offset, 0)
        max_y = min(max(bl[1], br[1]) + offset, hh - 1)
        return img_data[min_y:max_y, min_x:max_x, :], (min_x, min_y)


def get_cut_line(img, rows, cols, edge_removed):
    ybias, xbias = 0, 0
    if not edge_removed:
        img, bias = get_precut(img, rows, cols)
        xbias, ybias = bias
        print(xbias, ybias)
    hh, ww = img.shape[:2]
    #     if ww < 9000:
    col_lines, row_lines = get_cut_line_none_precut(img, rows, cols)
    #     else:
    #         col_lines, row_lines = get_cut_line_none_precut_edge(img, rows, cols)
    if xbias > 0:
        col_lines = [x + xbias for x in col_lines]
    if ybias > 0:
        row_lines = [x + ybias for x in row_lines]
    return col_lines, row_lines


def get_cut_line_none_precut(img, rows, cols):
    img_gray = img[:, :, 0]
    hh, ww = img.shape[:2]
    xstep = ww // cols
    ystep = hh // rows
    horizontal = np.sum(img_gray, axis=0)
    vertical = np.sum(img_gray, axis=1)
    # 电池片水平栅格线有可能干扰到电池片的水平切割，特别处理下
    conv_width_half = hh // (rows * 90)
    vertical = np.convolve(vertical,
                           np.ones(conv_width_half * 2 +
                                   1).astype('int'))[conv_width_half:-1 *
                                                     conv_width_half]
    ht = np.append([0], horizontal[:-1])
    hg = np.subtract(horizontal, ht)
    vg = np.diff(vertical)
    col_lines = [0]
    xstart = xstep * 3 // 4
    for x in range(cols - 1):
        xstart = np.argmax(hg[xstart:xstart + xstep // 2]) + xstart
        col_lines.append(xstart + 2)
        xstart += xstep * 3 // 4
    col_lines.append(ww - 1)

    row_lines = [0]
    ystart = ystep * 3 // 4
    for y in range(rows - 1):
        ystart = np.argmax(vg[ystart:ystart + ystep // 2]) + ystart
        row_lines.append(ystart + 2)
        ystart += ystep * 3 // 4
    row_lines.append(hh - 1)

    #verify
    # assert abs(col_lines[cols // 2] - ww // 2) / ww < 0.1, 'black edge'
    # assert abs(row_lines[rows // 2] - hh // 2) / hh < 0.1, 'black edge'

    return col_lines, row_lines


def get_cut_line_none_precut_edge(img, rows, cols):
    img_gray = img[:, :, 0]
    hh, ww = img.shape[:2]
    xstep = ww // cols
    ystep = hh // rows
    horizontal = np.sum(img_gray, axis=0)
    vertical = np.sum(img_gray, axis=1)
    # 电池片水平栅格线有可能干扰到电池片的水平切割，特别处理下
    conv_width_half = hh // (rows * 90)
    vertical = np.convolve(vertical,
                           np.ones(conv_width_half * 2 +
                                   1).astype('int'))[conv_width_half:-1 *
                                                     conv_width_half]
    ht = np.append([0], horizontal[:-1])
    hg = np.subtract(horizontal, ht)
    vg = np.diff(vertical)
    col_lines = [ww // 2]
    xform = ww // 2 - xstep * 3 // 4
    xlate = ww // 2 + xstep * 3 // 4
    for x in range(cols // 2):
        xform = (np.argmax(hg[max(xform - xstep // 2, 0):xform])) + xform - xstep // 2
        xlate = np.argmax(hg[xlate:xlate + xstep // 2]) + xlate
        col_lines.append(xform)
        col_lines.append(xlate)
        xform -= xstep * 3 // 4
        xlate += xstep * 3 // 4
    col_lines.sort()
    row_lines = [hh // 2]
    yform = hh // 2 - ystep * 3 // 4
    ylate = hh // 2 + ystep * 3 // 4
    for y in range(rows // 2):
        yform = (np.argmax(
            vg[max(yform - ystep // 2, 0):yform])) + yform - ystep // 2
        ylate = np.argmax(vg[ylate:ylate + ystep // 2]) + ylate
        row_lines.append(yform)
        row_lines.append(ylate)
        yform -= ystep * 3 // 4
        ylate += ystep * 3 // 4
    row_lines.sort()
    #verify
    # assert abs(col_lines[-1]-col_lines[-2]-xstep)/xstep < 0.4
    # assert abs(row_lines[-1]-row_lines[-2]-ystep)/ystep < 0.4

    return col_lines, row_lines


def grid_cell_lines(img_data, rows, cols, edge_removed):
    """
    切割线行列位置
    :param img_data: 图像数据
    :param rows: 电池片的行数
    :param cols: 电池片列数
    :param edge_removed: 是否无黑边
    :returns: 切割线位置
    """
    #     try:
    col_lines, row_lines = get_cut_line(img_data, rows, cols, edge_removed)
    #     except:
    #         logging.exception('wg cut failed, cut it by average.')
    #         stride_row, stride_col = int(img_data.shape[0]/rows), int(img_data.shape[1]/cols)
    #         row_lines=[x*stride_row for x in range(rows)]
    #         row_lines.append(img_data.shape[0])
    #         col_lines=[x*stride_col for x in range(cols)]
    #         col_lines.append(img_data.shape[1])

    return (row_lines, col_lines)


def collines_check(col_lines, cols, pad, ww):
    # 仅支持单行6片以上组件
    cellw = (col_lines[cols // 2 - 1] - col_lines[1] + col_lines[-2] -
             col_lines[cols // 2 + 1]) // (cols - 4) + pad
    col_lines[0] = max(col_lines[1] - cellw, 0)
    col_lines[-1] = min(col_lines[-2] + cellw, ww)
    mid_col_lines = col_lines[:cols // 2] + [
        col_lines[cols // 2 - 1] + cellw - pad
    ] + [col_lines[cols // 2 + 1] - cellw - pad] + col_lines[cols // 2 + 1:]
    return mid_col_lines


def rowlines_check(row_lines, rows, pad, hh):
    # 仅支持4行以上组件
    cellh = (row_lines[-2] - row_lines[1]) // (rows - 2) + pad
    row_lines[0] = max(row_lines[1] - cellh, 0)
    row_lines[-1] = min(row_lines[-2] + cellh, hh)
    return row_lines


def grid_cut(img_data, rows, cols, half_plate=True, edge_removed=False):
    """
    根据切线完成切图
    :param img_name: 图像名称
    :param img_data: 图像数据
    :param rows: 电池片的行数
    :param cols: 电池片列数
    :param half_plate: 是否半片
    :param edge_removed: 是否无黑边
    :param img_cut_folder: 保存路径
    :returns: 切割线位置
    """
    # cut_images = {}
    row_lines, col_lines = grid_cell_lines(img_data, rows, cols, edge_removed)
    pad = 0
    hh, ww = img_data.shape[:2]
    col_lines = collines_check(col_lines, cols, pad, ww)
    row_lines = rowlines_check(row_lines, rows, pad, hh)
    # step = 2 if half_plate else 1
    # for row in range(rows):
    #     for col in range(cols):
    #         if col in [cols // 2, cols // 2 - step + 1]:
    #             continue
    #         key = (row_lines[row], col_lines[col])
    #         cut_images[key] = img_data[row_lines[row]:row_lines[row + 1],
    #                                    col_lines[col]:col_lines[col + step], :]
    # col_lines_ret = col_lines[:cols // 2] + col_lines[cols // 2 + 1:]
    # print(row_lines, col_lines)
    return row_lines, col_lines